import asyncio
import websockets
import cv2
import base64
import json
from collections import Counter

cloud_uri = "ws://202.114.212.193:50511/detect"
front_video_uri = "ws://10.12.52.53:7979"
front_status_uri = "ws://10.21.188.130:7979"
window_size = 30
step = 10

def save_image(buffer, filename):
    with open(filename, "wb") as f:
        f.write(buffer)

async def capture_camera(queue1, queue2, frame_limit=10):
    cap = cv2.VideoCapture(0)
    try:
        while cap.isOpened() :
            ret, frame = cap.read()
            if not ret:
                print("Failed to read frame")
                break

            # 编码图片为 Base64
            _, buffer = cv2.imencode('.jpg', frame)
            # save_image(buffer, "test.jpg")
            img_str = base64.b64encode(buffer)
            await queue1.put(img_str)
            await queue2.put(img_str)
            await asyncio.sleep(0.01)  # 主动交出控制权，避免阻塞
    finally:
        cap.release()

async def video_websocket(URI, queue,):
    key = "cloud" if URI == cloud_uri else "front"
    window = []
    print(f"Connecting to {key}")
    try:
        async with websockets.connect(URI) as websocket:
            while True:
                img_str = await queue.get()  # 从队列中获取图片数据
                await websocket.send(img_str)
                print(f"Sent frame to server: {key}")

                if URI == cloud_uri:
                    # 接收服务器返回的推理结果
                    response = await websocket.recv()
                    result = json.loads(response)
                    window.append(result['detected_class'])
                    if len(window) == window_size:
                        counter = Counter(window)
                        most_common, _ = counter.most_common(1)[0]
                        # await queue.put(most_common)  # 将结果放入队列
                        # window = window[step:]
    except websockets.ConnectionClosed:
        print(f"Connection to {key} closed")
    except Exception as e:
        print(f"Failed to connect to {key} : {e}")

async def main():
    queue_cloud = asyncio.Queue(maxsize=10)
    queue_front = asyncio.Queue(maxsize=10)
    await asyncio.gather(
        capture_camera(queue_cloud, queue_front),
        video_websocket(cloud_uri, queue_cloud),
        video_websocket(front_video_uri, queue_front)
    )

asyncio.run(main())